#!/usr/bin/python3
import os, sys, json

#python3 extract_origin_data.py


# 用户列表文件 `user_list_GU2.txt`
user_file = './user_list_KKI2.txt'
# 源文件地址
source_data_path = '/home/li/Documents/MRI/ASD_data'
# 目标文件地址
target_data_path = '/mri_projects/ASD_Analysis/DataProcessed'

dir_t1 = 'T1'
dir_rest = 'rest'

def ext_target_dir(source_data_path, dir_t1, dir_rest, user_list, extract_dir = "GU", rest_sub_dir = "FunImg"):
    print('#' * 50)
    print('Start Checking File...')

    mk_dir_list = {}
    # 输出所有文件和文件夹
    for sub_dir in os.listdir(source_data_path):
        # 处理T1数据
        if sub_dir == dir_t1:
            input_dir = os.path.join(source_data_path, sub_dir)

            mk_dir_list[sub_dir] = []
            for user_sub_dir in os.listdir(input_dir):
                if user_sub_dir in user_list:
                    user_file_path = os.path.join(input_dir, user_sub_dir)
                    print('[EXTRACT DIR_NAME]: %s' % (user_file_path))
                    #
                    mk_dir_list[sub_dir].append(user_file_path)

        # 处理rest数据
        elif sub_dir == dir_rest:
            input_dir = os.path.join(source_data_path, sub_dir)
            input_dir = os.path.join(input_dir, extract_dir)
            input_dir = os.path.join(input_dir, rest_sub_dir)

            mk_dir_list[sub_dir] = []
            for user_sub_dir in os.listdir(input_dir):
                if user_sub_dir in user_list:
                    user_file_path = os.path.join(input_dir, user_sub_dir)
                    print('[EXTRACT DIR_NAME]: %s' % (user_file_path))
                    #
                    mk_dir_list[sub_dir].append(user_file_path)

    print("# Total T1 Files: %d" % (len(mk_dir_list[dir_t1])))
    print("# Total rest Files: %d" % (len(mk_dir_list[dir_rest])))
    print('Checking File Finished!')
    print('#' * 50)
    return mk_dir_list

def read_user_list(file_path):
    user_list = []

    fp = open(file_path, "r")
    for line in fp.readlines():
        line_list = line.split()
        item_name = line_list[0] + line_list[1] + '_' + line_list[3]

        user_list.append(item_name)

    return user_list

def rm_processed_file(dirs_dict, target_dirs, dir_t1, dir_rest, rm_t1 = True, rm_rest = True):
    target_t1_dir = os.path.join(target_dirs, dir_t1)
    target_rest_dir = os.path.join(target_dirs, dir_rest)

    print('#' * 50)
    print('Start Removing Subject...')

    for (dir_name, dir_list) in dirs_dict.items():
        file_list = [line_item.split('/')[-1] for line_item in dirs_dict[dir_t1]]
        if rm_t1 and dir_name == dir_t1:
            cmd_str = ""
            for dir_t1_item in file_list:
                target_t1_file = os.path.join(target_t1_dir, dir_t1_item)
                cmd_str = "rm -rf " + target_t1_file

                if os.path.exists(target_t1_file):
                    print("[Remove T1 Subject]: ==> %s" % (dir_t1_item))
                    print("[CMD]: %s" % (cmd_str))

                    os.system(cmd_str)
        if rm_rest and dir_name == dir_rest:
            cmd_str = ""
            for dir_rest_item in file_list:
                target_rest_file = os.path.join(target_rest_dir, dir_rest_item)
                cmd_str = "rm -rf " + target_rest_file

                if os.path.exists(target_rest_file):
                    print("[Remove rest Subject]: %s" % (dir_rest_item))
                    print("[CMD]: %s" % (cmd_str))

                    os.system(cmd_str)

    print('Remove Subject Finished!')
    print('#' * 50)

user_list = read_user_list(user_file)
dir_dict = ext_target_dir(source_data_path, dir_t1, dir_rest, user_list, "KKI")
rm_processed_file(dir_dict, target_data_path, dir_t1, dir_rest)
